"""
这个代码随机进行数据增强，对图片进行色彩增强、高斯噪声、水平镜像、放大、旋转、剪切

"""
import cv2
import os
from tqdm import tqdm
import numpy as np
import random
from random import randint
root_dir = "/media/liyan/3b75ab20-92d3-4be6-81c2-4e1798e2fe16/我的数据集/图像分类数据集/总/"
def contrast_brightness_image(src1:str, a, g, path_out):
    '''
        色彩增强（通过调节对比度和亮度）
    '''
    h, w, c = src1.shape  # 获取shape的数值，height和width、通道
    # 新建全零图片数组src2,将height和width，类型设置为原图片的通道类型(色素全为零，输出为全黑图片)
    src2 = np.zeros([h, w, c], src1.dtype)
    # addWeighted函数说明:计算两个图像阵列的加权和
    dst = cv2.addWeighted(src1, a, src2, 1 - a, g)
    cv2.imwrite(path_out, dst)

def gasuss_noise(image, path_out_gasuss, mean=0, var=0.001):
    '''
        添加高斯噪声
        mean : 均值
        var : 方差
    '''
    image = np.array(image / 255, dtype=float)
    noise = np.random.normal(mean, var ** 0.5, image.shape)
    out = image + noise
    if out.min() < 0:
        low_clip = -1.
    else:
        low_clip = 0.
    out = np.clip(out, low_clip, 1.0)
    out = np.uint8(out * 255)
    cv2.imwrite(path_out_gasuss, out)

def mirror(image, path_out_mirror):
    '''
        水平镜像
    '''
    h_flip = cv2.flip(image, 1)
    cv2.imwrite(path_out_mirror, h_flip)

def vertical_mirror(image,path_out_mirror):
    '''
        垂直镜像
    '''
    flip = cv2.flip(image,0)
    cv2.imwrite(path_out_mirror,flip)

def vertical_horizentor_mirror(image,path_out_mirror):
    '''
        水平垂直镜像
    '''
    flip = cv2.flip(image,-1)
    cv2.imwrite(path_out_mirror,flip)

def affine(image,path_out):
    '''
    仿射变换
    '''
    point1 = np.float32([[50, 50], [300, 50], [50, 200]])
    point2 = np.float32([[10, 100], [300, 50], [100, 250]])
    rows, cols = image.shape[:2]
    M = cv2.getAffineTransform(point1, point2)
    dst1 = cv2.warpAffine(image, M, (cols, rows), borderValue=(255, 255, 255))
    cv2.imwrite(path_out, dst1)

def resize(image, path_out_large):
    '''
        放大两倍
    '''
    height, width = image.shape[:2]
    large = cv2.resize(image, (2 * width, 2 * height))
    cv2.imwrite(path_out_large, large)

def rotate(image, path_out_rotate):
    '''
        旋转
    '''
    rows, cols = image.shape[:2]
    M = cv2.getRotationMatrix2D((cols / 2, rows / 2), 10, 1)
    dst = cv2.warpAffine(image, M, (cols, rows))
    cv2.imwrite(path_out_rotate, dst)


def shear(image, path_out_shear):
    '''
        剪切
    '''
    height, width = image.shape[:2]
    cropped = image[int(height / 9):height, int(width / 9):width]
    cv2.imwrite(path_out_shear, cropped)

if True:
    for data_path in os.listdir(root_dir):
        if os.path.isdir(root_dir+data_path):
            image_files = os.listdir(root_dir+data_path)
            images_files_length = len(image_files)
            if images_files_length < 3000:      #数据集足够不作增强
                supplement_length = 3000-images_files_length
                for i in tqdm(range(0, supplement_length),ncols=150):       #对不足的部分作增强
                    random_selected = random.choice(image_files)
                    image_path = root_dir+data_path+os.sep+random_selected
                    j = random.choice([0,1,2,3,4,5,6,7])
                    path_out = root_dir+data_path+os.sep+random_selected.split('.')[0]+f'{randint(0,1000)}.jpg'
                    image = cv2.imread(image_path)
                    if j == 0:
                        contrast_brightness_image(image, 1.2, 10, path_out)
                    elif j == 1:
                        gasuss_noise(image, path_out)
                    elif j == 2:
                        mirror(image, path_out)
                    elif j == 3:
                        resize(image, path_out)
                    elif j == 4:
                        rotate(image, path_out)
                    elif j == 5:
                        shear(image, path_out)
                    elif j == 6:
                        vertical_horizentor_mirror(image,path_out)
                    elif j == 7:
                        vertical_mirror(image,path_out)